"""
将所有的_finish.json文件提取出来整合成一个没有ID的表格，为后续做准备
"""
import json
import os
import csv
import cv2

# 定义一个列表，用于存储标注信息
train_data = []
val_data = []

# 遍历指定文件夹，找到目标JSON标注文件
for root, dirs, files in os.walk(r"D:\file\postgrad\experiment\bird_ava_dataset\videos_labelframes", topdown=False):
    files.sort()
    file_count = 0  # 计数文件的顺序，用于划分训练和验证

    for file in files:
        # 根据自己的命名换成新的后缀
        if "_finish_6.json" in file:
            file_count += 1
            jsonPath = os.path.join(root, file)
            with open(jsonPath, encoding='utf-8') as f:
                line = f.read()
                corrected_content = line.replace("'", '"')
                viaJson = json.loads(corrected_content)

                attributeNums = [0]
                for i in range(1, 10):
                    attribute_key = str(i)
                    if attribute_key in viaJson['attribute']:
                        cumulative_value = attributeNums[-1] + len(viaJson['attribute'][attribute_key]['options'])
                        attributeNums.append(cumulative_value)

                files = {file_id: file_info['fname'] for file_id, file_info in viaJson['file'].items()}
                sorted_files = dict(sorted(files.items(), key=lambda item: item[1]))

                for metadata in viaJson['metadata']:
                    imagen_x = viaJson['metadata'][metadata]
                    xy = imagen_x['xy'][1:]
                    vid = imagen_x['vid']
                    fname = sorted_files.get(vid)

                    if fname:
                        videoName = fname.split('_')[0]
                        frameId = int(fname.split('_')[1].split('.')[0])

                        for action in imagen_x['av']:
                            avs = imagen_x['av'][action]
                            if avs:
                                avArr = avs.split(',')
                                for av in avArr:
                                    imgPath = os.path.join(root, f"{videoName}_{str(frameId).zfill(6)}.jpg")
                                    imgTemp = cv2.imread(imgPath)

                                    if imgTemp is not None:
                                        sp = imgTemp.shape
                                        img_H, img_W = sp[0], sp[1]
                                        x1, y1 = xy[0] / img_W, xy[1] / img_H
                                        x2, y2 = (xy[0] + xy[2]) / img_W, (xy[1] + xy[3]) / img_H
                                        x1, x2 = max(0, min(1, x1)), max(0, min(1, x2))
                                        y1, y2 = max(0, min(1, y1)), max(0, min(1, y2))
                                        actionId = attributeNums[int(action) - 1] + int(av) + 1
                                        dict_entry = [videoName, frameId, x1, y1, x2, y2, actionId]

                                        # 每五个文件中，前四个文件划分为训练集，第五个文件划分为验证集
                                        if int(videoName) % 5 == 0:
                                            val_data.append(dict_entry)
                                        else:
                                            train_data.append(dict_entry)

# 按 videoName 和 frameId 对 train_data 和 val_data 进行排序
train_data.sort(key=lambda x: (int(x[0]), int(x[1])))
val_data.sort(key=lambda x: (int(x[0]), int(x[1])))

# 将 train_data 写入 train_without_personID.csv 文件
with open(r'D:\file\postgrad\experiment\bird_ava_dataset\train_without_personID.csv', "w", newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerows(train_data)

# 将 val_data 写入 val_without_personID.csv 文件
with open(r'D:\file\postgrad\experiment\bird_ava_dataset\val_without_personID.csv', "w", newline='') as csvfile:
    writer = csv.writer(csvfile)
    writer.writerows(val_data)


